"""
1.视频帧读取 → 灰度化 → 高斯模糊
2.Canny边缘检测 → 感兴趣区域（ROI）截取
3.霍夫变换检测直线 → 过滤无效线段
4.绘制车道线
"""
import cv2
import numpy as np


def img_preprocess(frame):
    # 1. 预处理
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blur, 50, 150)
    # cv2.imshow("edges", edges)
    # 2. 定义ROI（梯形区域）
    height, width = edges.shape
    return edges, height, width


def img_mask(edges, roi):
    mask = np.zeros_like(edges)
    cv2.fillPoly(mask, roi, 255)
    # cv2.imshow("mask", mask)
    masked_edges = cv2.bitwise_and(edges, mask)
    # cv2.imshow("masked_edges", masked_edges)
    return masked_edges


def line_detection(masked_edges):
    # 3. 霍夫变换检测直线
    lines = cv2.HoughLinesP(masked_edges, 1, np.pi / 180, threshold=50, minLineLength=100, maxLineGap=50)

    # print(lines)
    # 4. 绘制车道线
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            if y1 / y2 > 1.2 or y2 / y1 > 1.2:
                cv2.line(frame, (x1, y1), (x2, y2), (0, 0, 255), 5)
    return frame

    # cv2.imshow("Result", frame)
    # cv2.waitKey(0)


frame = cv2.imread("../images/chedaoxian.png")
frame = cv2.resize(frame, (0, 0), fx=0.3, fy=0.3)
img_edges, height, width = img_preprocess(frame)
ROI = np.array([[(0, height), (0, height // 2),
                 (width // 2 - 20, height // 3),
                 (width // 2 + 20, height // 3),
                 (width, height // 2), (width, height)]], dtype=np.int32)
masked_edges = img_mask(img_edges, ROI)
line_detection(masked_edges)
